﻿/*
    Copyright (c) 2012, iD Commerce + Logistics
    All rights reserved.

    Redistribution and use in source and binary forms, with or without modification, are permitted
    provided that the following conditions are met:

    Redistributions of source code must retain the above copyright notice, this list of conditions
    and the following disclaimer. Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions and the following disclaimer in the documentation and/or other
    materials provided with the distribution.
 
    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
    IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
    FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
    CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
    CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
    OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
    POSSIBILITY OF SUCH DAMAGE.
*/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Diagnostics;

namespace IdComLog.Data
{
    /// <summary>
    /// Helper class to generate IAsyncEnumerable state machines.
    /// </summary>
    public static class AsyncEnumerableEx
    {
        public static IAsyncEnumerable<T> Using<T, TState>(Func<CancellationToken, Task<TState>> createFunc, Func<TState, CancellationToken, Task<IAsyncEnumerable<T>>> enumerateFunc) where TState : IDisposable
        {
            return Using(createFunc, state => DeferAsync(ct => enumerateFunc(state, ct)));
        }

        public static IAsyncEnumerable<T> Using<T, TState>(Func<CancellationToken, Task<TState>> createFunc, Func<TState, IAsyncEnumerable<T>> enumerateFunc) where TState : IDisposable
        {
            return FromAsync(createFunc).SelectMany(x => AsyncEnumerable.Using(() => x, enumerateFunc));
        }

        public static IAsyncEnumerable<T> DeferAsync<T>(Func<CancellationToken, Task<IAsyncEnumerable<T>>> func)
        {
            return FromAsync(func).SelectMany(x => x);
        }

        public static IAsyncEnumerable<T> FromAsync<T>(Func<CancellationToken, Task<T>> func)
        {
            if (func == null) throw new ArgumentNullException("func");

            return new FromAsyncEnumerable<T>(func);
        }

        sealed class FromAsyncEnumerable<T> : IAsyncEnumerable<T>
        {
            readonly Func<CancellationToken, Task<T>> func;

            public FromAsyncEnumerable(Func<CancellationToken, Task<T>> func)
            {
                Debug.Assert(func != null);
                this.func = func;
            }

            public IAsyncEnumerator<T> GetEnumerator()
            {
                return new FromAsyncEnumerator<T>(func);
            }
        }

        sealed class FromAsyncEnumerator<T> : IAsyncEnumerator<T>
        {
            Func<CancellationToken, Task<T>> func;

            public FromAsyncEnumerator(Func<CancellationToken, Task<T>> func)
            {
                Debug.Assert(func != null);
                this.func = func;
            }

            public T Current { get; private set; }

            public async Task<bool> MoveNext(CancellationToken cancellationToken)
            {
                bool ret;

                if (ret = func != null)
                {
                    Current = await func(cancellationToken).ConfigureAwait(false);
                    func = null;
                }

                return ret;
            }

            public void Dispose()
            {
                func = null;
                Current = default(T);
            }
        }

        /// <summary>
        /// Creates an async enumerable from functional components.
        /// </summary>
        /// <typeparam name="T">The type returned by the enumerable.</typeparam>
        /// <typeparam name="TState">The state maintained for each iteration of the enumerable. If it implements IDisposable, it will be disposed.</typeparam>
        /// <param name="createFunc">A function to create the state, called on each iteration.</param>
        /// <param name="moveNextFunc">A function to move the iteration to the next value.</param>
        /// <param name="getValueFunc">A function to pull the current value from the state.</param>
        /// <returns>An async enumerable.</returns>
        public static IAsyncEnumerable<T> Create<T, TState>(Func<CancellationToken, Task<TState>> createFunc, Func<TState, CancellationToken, Task<bool>> moveNextFunc, Func<TState, T> getValueFunc)
        {
            if (createFunc == null) throw new ArgumentNullException("createFunc");
            if (moveNextFunc == null) throw new ArgumentNullException("moveNextFunc");
            if (getValueFunc == null) throw new ArgumentNullException("getValueFunc");

            return new FunctionalAsyncEnumerable<T, TState>(createFunc, moveNextFunc, getValueFunc);
        }

        /// <summary>
        /// Creates an async enumerable from functional components.
        /// </summary>
        /// <typeparam name="T">The type returned by the enumerable.</typeparam>
        /// <typeparam name="TState">The state maintained for each iteration of the enumerable. If it implements IDisposable, it will be disposed.</typeparam>
        /// <param name="createFunc">A function to create the state, called on each iteration.</param>
        /// <param name="moveNextFunc">A function to move the iteration to the next value.</param>
        /// <param name="getValueFunc">A function to pull the current value from the state.</param>
        /// <returns>An async enumerable.</returns>
        public static IAsyncEnumerable<T> Create<T, TState>(Func<TState> createFunc, Func<TState, CancellationToken, Task<bool>> moveNextFunc, Func<TState, T> getValueFunc)
        {
            if (createFunc == null) throw new ArgumentNullException("createFunc");

            return Create(tok => Task.FromResult(createFunc()), moveNextFunc, getValueFunc);
        }

        sealed class FunctionalAsyncEnumerable<T, TState> : IAsyncEnumerable<T>
        {
            readonly Func<CancellationToken, Task<TState>> createFunc;
            readonly Func<TState, CancellationToken, Task<bool>> moveNextFunc;
            readonly Func<TState, T> getValueFunc;

            public FunctionalAsyncEnumerable(Func<CancellationToken, Task<TState>> createFunc, Func<TState, CancellationToken, Task<bool>> moveNextFunc, Func<TState, T> getValueFunc)
            {
                Debug.Assert(createFunc != null);
                Debug.Assert(moveNextFunc != null);
                Debug.Assert(getValueFunc != null);

                this.createFunc = createFunc;
                this.moveNextFunc = moveNextFunc;
                this.getValueFunc = getValueFunc;
            }

            public IAsyncEnumerator<T> GetEnumerator()
            {
                return new FunctionalAsyncEnumerator<T, TState>(createFunc, moveNextFunc, getValueFunc);
            }
        }

        sealed class FunctionalAsyncEnumerator<T, TState> : IAsyncEnumerator<T>
        {
            Func<CancellationToken, Task<TState>> createFunc;
            Func<TState, CancellationToken, Task<bool>> moveNextFunc;
            Func<TState, T> getValueFunc;

            TState state;

            public T Current { get { return getValueFunc(state); } }

            public FunctionalAsyncEnumerator(Func<CancellationToken, Task<TState>> createFunc, Func<TState, CancellationToken, Task<bool>> moveNextFunc, Func<TState, T> getValueFunc)
            {
                Debug.Assert(createFunc != null);
                Debug.Assert(moveNextFunc != null);
                Debug.Assert(getValueFunc != null);

                this.createFunc = createFunc;
                this.moveNextFunc = moveNextFunc;
                this.getValueFunc = getValueFunc;
            }

            public Task<bool> MoveNext(CancellationToken cancellationToken)
            {
                if (createFunc != null)
                {
                    return MoveNextFirst(cancellationToken);
                }

                return moveNextFunc(state, cancellationToken);
            }

            async Task<bool> MoveNextFirst(CancellationToken cancellationToken)
            {
                state = await createFunc(cancellationToken).ConfigureAwait(false);
                createFunc = null;

                return await moveNextFunc(state, cancellationToken).ConfigureAwait(false);
            }

            public void Dispose()
            {
                if (createFunc == null && typeof(IDisposable).IsAssignableFrom(typeof(TState)))
                {
                    ((IDisposable)state).Dispose();
                }

                createFunc = null;
                moveNextFunc = null;
                getValueFunc = null;
                state = default(TState);
            }
        }

        /// <summary>
        /// Creates an IAsyncEnumerable with yield-like syntax.
        /// </summary>
        /// <typeparam name="T">The type returned by the enumerable.</typeparam>
        /// <param name="func">A funcion which yields a collection using the AsyncYielder.</param>
        /// <returns>An IAsyncEnumerable generated by <paramref name="func"/>.</returns>
        public static IAsyncEnumerable<T> Create<T>(Func<AsyncYielder<T>, CancellationToken, Task> func)
        {
            if (func == null) throw new ArgumentNullException("func");

            return new YieldAsyncEnumerable<T>(func);
        }

        /// <summary>
        /// An IAsyncEnumerable which creates YieldAsyncEnumerators.
        /// </summary>
        /// <typeparam name="T">The type returned by the enumerable.</typeparam>
        sealed class YieldAsyncEnumerable<T> : IAsyncEnumerable<T>
        {
            readonly Func<AsyncYielder<T>, CancellationToken, Task> func;

            public YieldAsyncEnumerable(Func<AsyncYielder<T>, CancellationToken, Task> func)
            {
                Debug.Assert(func != null);

                this.func = func;
            }

            public IAsyncEnumerator<T> GetEnumerator()
            {
                return new YieldAsyncEnumerator<T>(func);
            }
        }

        /// <summary>
        /// Enumerates over an async function.
        /// </summary>
        /// <typeparam name="T">The type returned by the enumerator.</typeparam>
        sealed class YieldAsyncEnumerator<T> : IAsyncEnumerator<T>
        {
            Func<AsyncYielder<T>, CancellationToken, Task> func;
            AsyncYielder<T> yielder;
            Task task;
            CancellationTokenSource cts;

            public YieldAsyncEnumerator(Func<AsyncYielder<T>, CancellationToken, Task> func)
            {
                Debug.Assert(func != null);

                this.func = func;
            }

            ~YieldAsyncEnumerator()
            {
                DisposeImpl();
            }

            public T Current { get; private set; }

            public async Task<bool> MoveNext(CancellationToken cancellationToken)
            {
                if (task != null)
                {
                    // Second MoveNext() call. Tell Yielder to let the function continue.
                    yielder.Continue();
                }
                else
                {
                    if (func == null)
                    {
                        throw new AsyncYielderDisposedException();
                    }

                    // First MoveNext() call. Start the task.

                    cts = new CancellationTokenSource();

                    yielder = new AsyncYielder<T>();

                    task = func(yielder, cts.Token);
                    func = null;
                }

                using (cancellationToken.Register(o => ((CancellationTokenSource)o).Cancel(), cts))
                {
                    // Wait for yield or return.

                    Task finished = await Task.WhenAny(task, yielder.YieldTask).ConfigureAwait(false);

                    if (finished != task)
                    {
                        // the function returned a result.

                        Current = yielder.YieldTask.Result;
                        return true;
                    }

                    // The operation is finished.

                    Task t = task;

                    yielder = null;
                    task = null;

                    if (t.IsFaulted)
                    {
                        throw t.Exception;
                    }

                    return false;
                }
            }

            public void Dispose()
            {
                DisposeImpl();
                GC.SuppressFinalize(this);
            }

            void DisposeImpl()
            {
                if (task != null)
                {
                    Task t = task;
                    yielder.Break();
                    yielder = null;
                    func = null;
                    try
                    {
                        t.Wait();
                    }
                    catch (AggregateException ex)
                    {
                        if (!(ex.GetBaseException() is AsyncYielderDisposedException))
                        {
                            throw;
                        }
                    }
                    finally
                    {
                        t.Dispose();
                        task = null;
                    }
                }
            }
        }
    }

    /// <summary>
    /// Controls yielding in an AsyncEnumerable.
    /// </summary>
    /// <typeparam name="T">The type to yield.</typeparam>
    public sealed class AsyncYielder<T>
    {
        TaskCompletionSource<T> setTcs = new TaskCompletionSource<T>();
        TaskCompletionSource<int> getTcs;

        internal Task<T> YieldTask { get { return setTcs.Task; } }

        /// <summary>
        /// Pushes a new value to be returned by the AsyncEnumerable.
        /// </summary>
        /// <param name="value">The value to return.</param>
        /// <returns>A Task representing the operation.</returns>
        public Task YieldReturn(T value)
        {
            getTcs = new TaskCompletionSource<int>();
            Task t = getTcs.Task;
            setTcs.SetResult(value);
            return t;
        }

        internal void Continue()
        {
            setTcs = new TaskCompletionSource<T>();
            getTcs.SetResult(0);
        }

        internal void Break()
        {
            AsyncYielderDisposedException ex = new AsyncYielderDisposedException();
            getTcs.TrySetException(ex);
            setTcs.TrySetException(ex);
        }
    }

    [Serializable]
    sealed class AsyncYielderDisposedException : ObjectDisposedException
    {
        internal AsyncYielderDisposedException()
            : base("AsyncYielder")
        {
        }
    }
}
